import tkinter as tk
from tkinter import ttk, messagebox
from PIL import Image, ImageTk
import cv2
import numpy as np
import math
import mediapipe as mp
import joblib
import time


class App:
    def __init__(self, root):
        self.root = root
        self.root.title("# OpenCV 实时检测")
        self.root.geometry("800x600")  # 设置窗口大小为800x600

        # 创建显示OpenCV结果的标签
        self.label = tk.Label(root, width=640, height=480)
        self.label.pack(pady=20)

        # 创建按钮和下拉菜单框架
        self.control_frame = tk.Frame(root)
        self.control_frame.pack(pady=10)

        self.start_btn = tk.Button(self.control_frame, text="开始", command=self.start)
        self.start_btn.grid(row=0, column=0, padx=5)

        self.stop_btn = tk.Button(self.control_frame, text="停止", command=self.stop)
        self.stop_btn.grid(row=0, column=1, padx=5)

        self.model_var = tk.StringVar()
        self.model_var.set("选择模型")

        self.model_menu = ttk.Combobox(self.control_frame, textvariable=self.model_var)
        self.model_menu['values'] = ("数字检测", "石头剪刀布")
        self.model_menu.grid(row=0, column=2, padx=5)

        self.cap = None
        self.running = False
        self.selected_model = None

        self.show_grey_image()  # 初始显示灰色填充

        # 设置姿态检测和手部检测
        self.mpPose = mp.solutions.pose
        self.pose = self.mpPose.Pose(static_image_mode=False, smooth_landmarks=True, min_detection_confidence=0.5,
                                     min_tracking_confidence=0.5)
        self.mphands = mp.solutions.hands
        self.hands = self.mphands.Hands(static_image_mode=False, max_num_hands=2, min_detection_confidence=0.75,
                                        min_tracking_confidence=0.75)
        self.mpDraw = mp.solutions.drawing_utils

        self.pTime = 0

    def start(self):
        if self.cap is None:
            self.cap = cv2.VideoCapture(0)  # 打开默认相机
        self.selected_model = self.model_var.get()
        if self.selected_model not in ["数字检测", "石头剪刀布"]:
            messagebox.showerror("错误", "请选择有效的模型")
            return
        self.running = True
        self.update_frame()

    def stop(self):
        self.running = False
        self.label.config(image='')  # 清除显示
        self.show_grey_image()

    def update_frame(self):
        if self.running:
            ret, img = self.cap.read()
            if ret:
                imgRGB = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
                results = self.pose.process(imgRGB)
                hand = self.hands.process(imgRGB)

                if results.pose_landmarks:
                    self.mpDraw.draw_landmarks(img, results.pose_landmarks, self.mpPose.POSE_CONNECTIONS)

                res_hand = []
                pix = []
                if hand.multi_hand_landmarks:
                    for hand_landmarks in hand.multi_hand_landmarks:
                        for item in hand_landmarks.landmark:
                            res_hand.append([item.x, item.y])
                        self.mpDraw.draw_landmarks(img, hand_landmarks, self.mphands.HAND_CONNECTIONS)

                image_w, image_h, _ = img.shape
                for xy in res_hand:
                    pix.append([xy[0] * image_w, xy[1] * image_h])

                dis = []
                cont = ''
                label = ''
                if pix:
                    for i in range(20):
                        dis.append(math.pow(pix[i + 1][0] - pix[0][0], 2) + math.pow(pix[i + 1][1] - pix[0][1], 2))
                    predata = np.array(dis)
                    predata = predata.reshape(1, -1)
                    if self.selected_model == '数字检测':
                        knn = joblib.load('num_classifier.pkl')
                        labels = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']
                        cont = 'Number:  '
                    elif self.selected_model == '石头剪刀布':
                        knn = joblib.load('knn_classifier.pkl')
                        labels = ['scissor','stone','cloth']
                        cont = 'Pose:  '
                    predictions = knn.predict(predata)
                    for i, prediction in enumerate(predictions):
                        label = labels[prediction]

                cTime = time.time()
                fps = 1 / (cTime - self.pTime)
                self.pTime = cTime

                cv2.putText(img, 'FPS:  '+str(int(fps)), (10, 50), cv2.FONT_HERSHEY_PLAIN, 2, (0, 0, 255), 3)
                cv2.putText(img, cont + label, (400, 50), cv2.FONT_HERSHEY_PLAIN, 2, (0, 0, 255), 3)

                cv2image = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
                img = Image.fromarray(cv2image)
                imgtk = ImageTk.PhotoImage(image=img)
                self.label.imgtk = imgtk
                self.label.config(image=imgtk)
            self.root.after(10, self.update_frame)

    def show_grey_image(self):
        grey_img = np.zeros((480, 640, 3), dtype=np.uint8) + 128  # 创建灰色图像
        img = Image.fromarray(grey_img)
        imgtk = ImageTk.PhotoImage(image=img)
        self.label.imgtk = imgtk
        self.label.config(image=imgtk)

    def on_closing(self):
        self.running = False
        if self.cap:
            self.cap.release()
        self.root.destroy()


if __name__ == "__main__":
    root = tk.Tk()
    app = App(root)
    root.protocol("WM_DELETE_WINDOW", app.on_closing)
    root.mainloop()
